# 导入测试基类，实现自动初始化
from tests.base_test import base_test_case

# 文本识别demo
import json
import os

from paddleocr import TextRecognition, PaddleOCR

current_dir = os.path.dirname(os.path.abspath(__file__))


def text_recognition_demo(image_path):
    """
    使用TextRecognition进行文本识别
    """
    print("使用TextRecognition进行文本识别")
    # 1. 创建TextRecognition实例
    # 默认使用PP-OCRv5_server_rec模型
    recognizer = TextRecognition(
        model_name="PP-OCRv5_server_rec",  # 可选，使用默认模型
    )

    # 2. 准备输入图像

    # 3. 执行文本识别
    results = recognizer.predict(input=image_path)

    # 4. 处理识别结果
    for result in results:
        # 打印识别结果
        result.print()

        # # 保存结果到图像
        # result.save_to_img(save_path="./output/")
        #
        # # 保存结果到JSON文件
        # result.save_to_json(save_path="./output/recognition_result.json")

        # 直接访问识别文本和置信度
        # print(f"识别文本: {result.text}")
        # print(f"置信度: {result.score}")


def text_recognition_with_custom_model(image_path):
    """
    使用自定义模型的文本识别demo
    """
    print("使用移动端模型进行文本识别")
    # 使用不同的预训练模型
    recognizer = TextRecognition(
        model_name="PP-OCRv5_mobile_rec"  # 使用移动端模型
    )

    results = recognizer.predict(input=image_path)

    for result in results:
        result.print()


def text_recognition_with_character_coordinates(image_path):
    """
    使用PaddleOCR进行文本识别并获取字符级坐标信息
    """
    print("使用PaddleOCR进行文本识别并获取字符级坐标信息")
    # 创建PaddleOCR实例，启用字符级坐标返回
    ocr = PaddleOCR(
        text_recognition_model_name="PP-OCRv5_server_rec",
        text_detection_model_name=None,
        use_textline_orientation=False,
        use_doc_orientation_classify=False,  # 禁用文档预处理
        use_doc_unwarping=False,  # 禁用文档矫正
        return_word_box=True  # 返回单词/字符坐标
    )

    # 执行OCR识别
    results = ocr.predict(input=image_path)

    # 处理结果
    for result in results:
        # 打印完整结果
        # 获取JSON数据并手动格式化
        # formatted_json = json.dumps(result.json, indent=2, ensure_ascii=False)
        # print(formatted_json)

        # 通过字典方式访问属性（根据日志输出分析）
        res_data = result.json['res']
        # print(f"检测到的文本框坐标: {res_data['dt_polys']}")
        # print(f"识别文本: {res_data['rec_texts']}")
        # print(f"置信度: {res_data['rec_scores']}")
        # print(f"文本类型: {res_data['text_type']}")

        # 如果启用了return_word_box，获取字符级坐标
        if 'text_word_boxes' in res_data and 'text_word' in res_data:
            print("字符级坐标:")
            for i, (word_boxes, words) in enumerate(zip(res_data['text_word_boxes'], res_data['text_word'])):
                print(f"  文本行 {i}:")
                for j, (box, char) in enumerate(zip(word_boxes, words)):
                    print(f"    字符 '{char}' 坐标: {box}")
    format_ocr_results_as_json(results)


def format_ocr_results_as_json(results):
    """
    将PaddleOCR的识别结果格式化为JSON数据
    """
    # 按指定格式处理结果
    for result in results:
        res_data = result.json['res']

        # 处理每个检测到的文本框
        dt_polys = res_data.get('dt_polys', [])
        rec_texts = res_data.get('rec_texts', [])
        rec_scores = res_data.get('rec_scores', [])

        # 确保三个列表长度一致
        min_length = min(len(dt_polys), len(rec_texts), len(rec_scores))

        for i in range(min_length):
            # 构造符合要求的JSON格式
            formatted_result = {
                "text": rec_texts[i],
                "bbox": dt_polys[i],
                "score": rec_scores[i]
            }

            # 输出格式化的JSON字符串
            print(json.dumps(formatted_result, ensure_ascii=False))


if __name__ == "__main__":
    image_path = os.path.join(current_dir, "test_files", "textline.png")
    # image_path = os.path.join(current_dir, "test_files", "book.jpg")
    # image_path = os.path.join(current_dir, "test_files", "formula.png")
    # image_path = os.path.join(current_dir, "test_files", "table.jpg")

    # 运行文本识别demo
    # text_recognition_demo(image_path)

    # 运行自定义模型demo
    # text_recognition_with_custom_model(image_path)
    text_recognition_with_character_coordinates(image_path)
